"""
项目WebSocket服务
扩展现有WebSocket服务以支持项目房间和协同编辑
"""

import json
import asyncio
from typing import Dict, List, Any, Optional
from datetime import datetime
from fastapi import WebSocket, WebSocketDisconnect
from src.utils.logging import get_logger
from src.services.websocket_service import ConnectionManager
from src.services.enhanced_collaboration_service import get_enhanced_collaboration_service, Operation

logger = get_logger(__name__)

# 获取增强版协同编辑服务实例
enhanced_service = get_enhanced_collaboration_service()

class ProjectConnectionManager(ConnectionManager):
    """项目专用的WebSocket连接管理器"""
    
    def __init__(self):
        super().__init__()
        # 存储项目文档状态: {project_id: {document_id: document_content}}
        self.project_documents: Dict[str, Dict[str, Any]] = {}
        # 存储协同编辑操作: {project_id: [operation_history]}
        self.collaboration_operations: Dict[str, List[Dict[str, Any]]] = {}
        # 存储用户光标位置: {project_id: {user_id: cursor_position}}
        self.user_cursors: Dict[str, Dict[str, Dict[str, Any]]] = {}
    
    async def handle_document_operation(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理文档协同编辑操作"""
        document_id = message.get("document_id", "")
        operation_data = message.get("operation", {})
        
        if not document_id:
            return
        
        # 使用增强版协同编辑服务处理操作
        full_document_id = f"{project_id}_{document_id}"
        await enhanced_service.handle_document_operation(message, full_document_id, user_id)
    
    async def broadcast_to_project(self, message: Dict[str, Any], project_id: str, exclude_user: Optional[str] = None):
        """向项目中的所有用户广播消息"""
        # 项目房间使用格式: "project_{project_id}"
        session_id = f"project_{project_id}"
        await super().broadcast_to_session(message, session_id, exclude_user)
    
    def update_user_cursor(self, project_id: str, user_id: str, cursor_data: Dict[str, Any]):
        """更新用户光标位置"""
        if project_id not in self.user_cursors:
            self.user_cursors[project_id] = {}
        self.user_cursors[project_id][user_id] = cursor_data
    
    def get_project_document(self, project_id: str, document_id: str) -> Optional[str]:
        """获取项目文档内容"""
        full_document_id = f"{project_id}_{document_id}"
        return enhanced_service.manager.get_document_content(full_document_id)
    
    def get_collaboration_history(self, project_id: str, limit: int = 50) -> List[Dict[str, Any]]:
        """获取协作操作历史"""
        history = self.collaboration_operations.get(project_id, [])
        return history[-limit:] if history else []

# 扩展WebSocket服务以支持项目管理
class ProjectWebSocketService:
    def __init__(self):
        self.connection_manager = ProjectConnectionManager()
    
    async def handle_project_connection(self, websocket: WebSocket, project_id: str, user_id: str, user_info: Dict[str, Any]):
        """处理项目WebSocket连接"""
        session_id = f"project_{project_id}"
        await self.connection_manager.connect(websocket, session_id, user_id, user_info)
        
        try:
            while True:
                data = await websocket.receive_text()
                message = json.loads(data)
                await self.process_project_message(message, project_id, user_id)
                
        except WebSocketDisconnect:
            logger.info(f"项目WebSocket连接断开: 用户 {user_id}, 项目 {project_id}")
        finally:
            self.connection_manager.disconnect(session_id, user_id)
    
    async def process_project_message(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理项目相关消息"""
        message_type = message.get("type")
        
        if message_type == "document_operation":
            await self.connection_manager.handle_document_operation(message, project_id, user_id)
        elif message_type == "cursor_update":
            await self.handle_cursor_update(message, project_id, user_id)
        elif message_type == "project_update":
            await self.handle_project_update(message, project_id, user_id)
        elif message_type == "role_assignment":
            await self.handle_role_assignment(message, project_id, user_id)
        elif message_type == "feedback":
            await self.handle_feedback(message, project_id, user_id)
        elif message_type == "document_lock":
            await self.handle_document_lock(message, project_id, user_id)
        elif message_type == "document_unlock":
            await self.handle_document_unlock(message, project_id, user_id)
        elif message_type == "sync_request":
            await self.handle_sync_request(message, project_id, user_id)
        else:
            # 转发到基础WebSocket服务
            session_id = f"project_{project_id}"
            await self.connection_manager.broadcast_to_session(message, session_id, exclude_user=user_id)
    
    async def handle_cursor_update(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理光标位置更新"""
        cursor_data = message.get("cursor", {})
        self.connection_manager.update_user_cursor(project_id, user_id, cursor_data)
        
        # 广播光标位置给其他用户
        await self.connection_manager.broadcast_to_project(
            {
                "type": "cursor_update",
                "user_id": user_id,
                "cursor": cursor_data,
                "timestamp": datetime.now().isoformat()
            },
            project_id,
            exclude_user=user_id
        )
    
    async def handle_project_update(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理项目状态更新"""
        update_data = message.get("update", {})
        
        # 广播项目更新
        await self.connection_manager.broadcast_to_project(
            {
                "type": "project_update",
                "user_id": user_id,
                "update": update_data,
                "timestamp": datetime.now().isoformat()
            },
            project_id
        )
    
    async def handle_role_assignment(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理角色分配"""
        assignments = message.get("assignments", [])
        
        # 广播角色分配
        await self.connection_manager.broadcast_to_project(
            {
                "type": "role_assignment",
                "user_id": user_id,
                "assignments": assignments,
                "timestamp": datetime.now().isoformat()
            },
            project_id
        )
    
    async def handle_feedback(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理团队反馈"""
        feedback_data = message.get("feedback", {})
        
        # 广播反馈消息
        await self.connection_manager.broadcast_to_project(
            {
                "type": "feedback",
                "user_id": user_id,
                "feedback": feedback_data,
                "timestamp": datetime.now().isoformat()
            },
            project_id
        )
    
    async def handle_document_lock(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理文档锁定请求"""
        document_id = message.get("document_id", "")
        lock_info = message.get("lock_info", {})
        
        if document_id:
            full_document_id = f"{project_id}_{document_id}"
            await enhanced_service.handle_lock_request(
                {"lock_info": lock_info, "document_id": full_document_id}, 
                full_document_id, 
                user_id
            )
    
    async def handle_document_unlock(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理文档解锁请求"""
        document_id = message.get("document_id", "")
        
        if document_id:
            full_document_id = f"{project_id}_{document_id}"
            await enhanced_service.handle_unlock_request(
                {"document_id": full_document_id}, 
                full_document_id, 
                user_id
            )
    
    async def handle_sync_request(self, message: Dict[str, Any], project_id: str, user_id: str):
        """处理同步请求"""
        document_id = message.get("document_id", "")
        
        if document_id:
            full_document_id = f"{project_id}_{document_id}"
            await enhanced_service.handle_sync_request(
                {"document_id": full_document_id}, 
                full_document_id, 
                user_id
            )
    
    def get_project_users(self, project_id: str) -> Dict[str, Dict[str, Any]]:
        """获取项目用户列表"""
        session_id = f"project_{project_id}"
        return self.connection_manager.get_session_users(session_id)
    
    def get_collaboration_history(self, project_id: str, limit: int = 50) -> List[Dict[str, Any]]:
        """获取协作历史"""
        return self.connection_manager.get_collaboration_history(project_id, limit)


# 全局实例
project_websocket_service = ProjectWebSocketService()


def get_project_websocket_service() -> ProjectWebSocketService:
    """获取项目WebSocket服务实例"""
    return project_websocket_service